package com.example.demo.Service;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;

@Service
@Slf4j
public class WebSocketService extends TextWebSocketHandler {

    // 存储用户会话的 Map，键是用户 ID，值是一个 WebSocket 会话列表
    private final Map<Integer, List<WebSocketSession>> userSessions = new ConcurrentHashMap<>();

    // 用户连接时保存会话
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        Integer userId = Integer.valueOf(session.getUri().getPath().split("/")[2]);

        // 获取用户的现有连接列表，如果没有，创建新的
        List<WebSocketSession> sessions = userSessions.computeIfAbsent(userId, k -> new CopyOnWriteArrayList<>());

        // 添加新的 WebSocket 会话
        sessions.add(session);
        log.info("用户 {} 连接成功, 当前连接数: {}", userId, sessions.size());
    }

    // 处理收到的消息
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        log.info("收到用户 {} 的消息: {}", session.getId(), message.getPayload());
        session.sendMessage(new TextMessage("消息已接收"));
    }

    // 用户断开连接时移除会话
    @Override
    public void afterConnectionClosed(WebSocketSession session, org.springframework.web.socket.CloseStatus status) throws Exception {
        Integer userId = Integer.valueOf(session.getUri().getPath().split("/")[2]);

        // 从会话列表中移除该会话
        List<WebSocketSession> sessions = userSessions.get(userId);
        if (sessions != null) {
            sessions.remove(session);
            log.info("用户 {} 连接关闭, 当前连接数: {}", userId, sessions.size());
        }
    }

    // 发送通知消息给指定用户
    public void sendToUser(Integer userId, String message) {
        List<WebSocketSession> sessions = userSessions.get(userId);
        if (sessions != null) {
            for (WebSocketSession session : sessions) {
                if (session.isOpen()) {
                    try {
                        session.sendMessage(new TextMessage(message));
                        log.info("向用户 {} 发送消息: {}", userId, message);
                    } catch (IOException e) {
                        log.error("向用户 {} 发送消息失败: {}", userId, e.getMessage());
                    }
                }
            }
        } else {
            log.error("无法发送消息，用户 {} 会话不存在", userId);
        }
    }
}
